#pragma once

#include "macros.h"
#include <string>
#include <vector>
#include <opencv2/opencv.hpp>

cv::Mat preprocess_img(cv::Mat& img, int input_w, int input_h);

//! \class Int8EntropyCalibrator2
//!
//! \brief Implements Entropy calibrator 2.
//!  CalibrationAlgoType is kENTROPY_CALIBRATION_2.
//!
class Int8EntropyCalibrator2 : public nvinfer1::IInt8EntropyCalibrator2 {
 public:
  Int8EntropyCalibrator2(int batchsize, int input_w, int input_h, const char* img_dir, const char* calib_table_name, const char* input_blob_name, bool read_cache = true);

  virtual ~Int8EntropyCalibrator2();
  int getBatchSize() const TRT_NOEXCEPT override;
  bool getBatch(void* bindings[], const char* names[], int nbBindings) TRT_NOEXCEPT override;
  const void* readCalibrationCache(size_t& length) TRT_NOEXCEPT override;
  void writeCalibrationCache(const void* cache, size_t length) TRT_NOEXCEPT override;

 private:
  int batchsize_;
  int input_w_;
  int input_h_;
  int img_idx_;
  std::string img_dir_;
  std::vector<std::string> img_files_;
  size_t input_count_;
  std::string calib_table_name_;
  const char* input_blob_name_;
  bool read_cache_;
  void* device_input_;
  std::vector<char> calib_cache_;
};

